import os
import sys
import tree
import torch
from torch.utils.data import Dataset, DataLoader


root_dir = os.path.dirname(os.path.abspath(__file__))
print(root_dir)
sys.path.append(root_dir)
from GVPmodel import GVPTransCond, geo_batch
from inference_ribodiffusion import get_config
from RNADataset import RNADatasetV3, collate_fn, graph_collate_fn



def predict_bak():
    config = get_config()
    model = GVPTransCond(config)
    # print(model)
    dataset = RNADatasetV3(
        data_path="/data/slz/sais_medicine/saisdata",
        is_train=True
    )
    for data in dataset:
        # 所有变量都增加一个batch维度
        struct_data = tree.map_structure(lambda x:
            x.unsqueeze(0).repeat_interleave(1, dim=0),
            data)
        batch, batch_size, length = geo_batch(struct_data)
        output = model.struct_forward(batch, batch_size, length)
        # output = model.forward(struct_data)
        print(output.shape)
        break
    
def predict():
    config = get_config()
    model = GVPTransCond(config)
    # print(model)
    dataset = RNADatasetV3(
        data_path="/data/slz/sais_medicine/saisdata",
        is_train=True
    )
    dataloader = DataLoader(dataset, batch_size=2, collate_fn=collate_fn)
    for data in dataloader:
        # 所有变量都增加一个batch维度
        # struct_data = tree.map_structure(lambda x:
        #     x.unsqueeze(0).repeat_interleave(1, dim=0),
        #     data)
        output = model.forward(data)
        print(output.shape)
        break

if __name__ == "__main__":
    predict()
    # predict_bak()

    
    
    